#!/usr/bin/env python  
#-*- coding:utf-8 _*-  
""" 
@author:hello_life 
@license: Apache Licence 
@file: data.py 
@time: 2022/03/27
@software: PyCharm 
description:
"""
import os,sys
sys.path.insert(0,os.path.dirname(os.getcwd()))


import torch
import pandas as pd

from torch.utils.data import Dataset


from .parameters import Config
from utils.data_utils import content_to_id,label_to_id

config=Config()

def load_mnist(path, kind='train'):
    import os
    import gzip
    import numpy as np

    """Load MNIST data from `path`"""
    labels_path = os.path.join(path,
                               '%s-labels-idx1-ubyte.gz'
                               % kind)
    images_path = os.path.join(path,
                               '%s-images-idx3-ubyte.gz'
                               % kind)

    with gzip.open(labels_path, 'rb') as lbpath:
        labels = np.frombuffer(lbpath.read(), dtype=np.uint8,
                               offset=8)

    with gzip.open(images_path, 'rb') as imgpath:
        images = np.frombuffer(imgpath.read(), dtype=np.uint8,
                               offset=16).reshape(len(labels), 784)

    return images, labels

class CustomImageDataset(Dataset):
    def __init__(self,img_dir,kind,transform=True,target_transform=True):
        self.imgs,self.img_labels=load_mnist(img_dir,kind)
        self.do_transform=transform
        self.do_target_transform=target_transform

    def __len__(self):
        assert len(self.imgs)==len(self.img_labels)
        return len(self.img_labels)

    def __getitem__(self, idx):
        image,label=self.imgs[idx],self.img_labels[idx],
        if self.do_transform:
            image=self.transform(image)
        if self.target_transform:
            label=self.target_transform(label)
        return image,label

    def transform(self,x):
        x=torch.tensor(x,dtype=torch.float)
        return x

    def target_transform(self,x):
        x=torch.tensor(x,dtype=torch.long)
        return x

class IMDBDataset(Dataset):
    def __init__(self,config):
        self.config=config
        self.data=pd.read_csv(self.config.data_path)

    def __len__(self):
        assert len(self.data["review"])==len(self.data["sentiment"])
        return len(self.data["review"])

    def __getitem__(self, idx):
        x,y=self.data["review"].iloc[idx],self.data["sentiment"].iloc[idx]
        x=self.x_transformer(x)
        y=self.y_transformer(y)
        return x,y

    def x_transformer(self,x):
        x=content_to_id(x,self.config)
        return torch.tensor(x,dtype=torch.int32).to(self.config.device)

    def y_transformer(self,y):
        y=label_to_id(y)
        return torch.tensor(y,dtype=torch.int64).to(self.config.device)

